# -*- coding: utf-8 -*-
"""
    @project: pythonProject
    @Author：
    @file： fftmat.py
    @date：2025/3/19 14:17
    @blogs:
"""

import math
import numpy as np
import matplotlib.pyplot as plt
from scipy.io import loadmat

# 从 .mat 文件中读取信号数据
def load_signal_from_mat(file_path, signal_key='signal'):
    """
    从 .mat 文件中加载信号数据。
    :param file_path: .mat 文件路径
    :param signal_key: 信号在 .mat 文件中的键名
    :return: 信号数据 (numpy array)
    """
    data = loadmat(file_path)  # 加载 .mat 文件
    if signal_key not in data:
        raise KeyError(f"'{signal_key}' not found in the .mat file.")
    signal = data[signal_key].flatten()  # 将信号展平为一维数组
    freq = data['XDelta'].flatten()
    return signal,freq[0]

# 计算信号功率
def calculate_power(signal, fs, window='hann', n_fft=None):
    """
    计算信号的功率，考虑频谱泄露和栅栏效应。
    :param signal: 输入信号 (numpy array)
    :param fs: 采样频率 (Hz)
    :param window: 窗函数类型 ('hann', 'hamming', 等)
    :param n_fft: FFT点数 (可选，默认为信号长度)
    :return: 总功率 (float), 频率轴 (numpy array), 功率谱 (numpy array)
    """
    N = len(signal)  # 信号长度
    if n_fft is None:
        n_fft = N  # 默认FFT点数为信号长度
    elif n_fft < N:
        raise ValueError("n_fft must be greater than or equal to the signal length.")

    # 加窗以减少频谱泄露
    if window == 'hann':
        window = np.hanning(N)
    elif window == 'hamming':
        window = np.hamming(N)
    else:
        raise ValueError("Unsupported window type.")
    windowed_signal = signal * window

    # 零填充以减少栅栏效应
    fft_result = np.fft.fft(windowed_signal, n=n_fft)

    # 计算功率谱
    power_spectrum = (np.abs(fft_result) ** 2) / (n_fft * np.sum(window ** 2))/100  # 功率谱修正

    # 计算总功率（所有频率分量的功率之和）
    total_power = np.sum(power_spectrum)  # 总功率

    # 正频率部分的频率轴和功率谱
    positive_freq = np.fft.fftfreq(n_fft, d=1/fs)[:n_fft // 2]  # 正频率轴
    positive_power_spectrum = power_spectrum[:n_fft // 2]  # 正频率功率谱

    return total_power, positive_freq, positive_power_spectrum

def find_max_index(arr):
    np_array = np.array(arr)
    return np.max(np_array), np.argmax(np_array)

def calc_compensation_max_value(fft_mag, peak_idx, N) :
    if 1 <= peak_idx < N - 1:
        y0, y1, y2 = fft_mag[peak_idx - 1:peak_idx + 2]
        # 峰值幅度校正（减少栅栏效应影响）
        dy = y0 - y2
        denominator = 4 * y1 - 2 * y0 - 2 * y2
        if denominator != 0:  # 避免除零错误
            corrected_peak_mag = y1 - 0.25 * dy ** 2 / denominator
        else:
            corrected_peak_mag = y1
    else:
        corrected_peak_mag = fft_mag[peak_idx]
    return corrected_peak_mag

# 计算信号功率
def calculate_power2(signal, fs, window='hann', n_fft=None):
    """
    计算信号的功率，考虑频谱泄露和栅栏效应。
    :param signal: 输入信号 (numpy array)
    :param fs: 采样频率 (Hz)
    :param window: 窗函数类型 ('hann', 'hamming', 等)
    :param n_fft: FFT点数 (可选，默认为信号长度)
    :return: 总功率 (float), 频率轴 (numpy array), 功率谱 (numpy array)
    """
    N = len(signal)  # 信号长度
    if n_fft is None:
        n_fft = N  # 默认FFT点数为信号长度
    elif n_fft < N:
        raise ValueError("n_fft must be greater than or equal to the signal length.")

    # 加窗以减少频谱泄露
    if window == 'hann':
        window = np.hanning(N)
    elif window == 'hamming':
        window = np.hamming(N)
    else:
        raise ValueError("Unsupported window type.")
    windowed_signal = signal * window

    # 零填充以减少栅栏效应
    fft_result = np.fft.fft(windowed_signal, n=n_fft)

    # 计算功率谱
    power_spectrum = (np.abs(fft_result) ** 2) / (n_fft * np.sum(window ** 2))/100  # 功率谱修正

    max, max_indx = find_max_index(power_spectrum);

    # 计算总功率（所有频率分量的功率之和）
    #total_power = np.sum(power_spectrum[max_indx-130 : max_indx+130])  # 总功率
    total_power = power_spectrum[max_indx]  # 总功率

    corrected_peak_mag = calc_compensation_max_value(power_spectrum, max_indx, n_fft);

    # 正频率部分的频率轴和功率谱
    positive_freq = np.fft.fftfreq(n_fft, d=1/fs)[:n_fft // 2]  # 正频率轴
    positive_power_spectrum = power_spectrum[:n_fft // 2]  # 正频率功率谱

    return total_power, positive_freq, positive_power_spectrum

# 主程序
if __name__ == "__main__":
    # 从 .mat 文件中加载信号
    # file_path = '8_182397.mat'  # .mat 文件路径
    file_path = '45005798_34_483425.mat'  # .mat 文件路径
    # file_path = '15136718_22_078682.mat'  # .mat 文件路径

    signal_key = 'Y'  # 信号在 .mat 文件中的键名
    signal,freq = load_signal_from_mat(file_path, signal_key)

    # 设置采样频率
    fs = 1/freq  # 采样频率 (Hz)

    # 计算信号功率
    # total_power, freq, power_spectrum = calculate_power2(signal, fs, window='hann', n_fft=2 * len(signal))
    total_power, freq, power_spectrum = calculate_power2(signal, fs, window='hann', n_fft= 2 * len(signal))

    # 输出结果
    print("总功率:", total_power)

    # 输出结果(dBm)
    total_dbm = 10 * math.log(total_power, 10) + 30
    print("总功率:", total_dbm)

    # 绘制功率谱
    plt.figure(figsize=(10, 6))
    plt.plot(freq, power_spectrum)
    plt.title("Power Spectrum (with Hanning Window and Zero-Padding)")
    plt.xlabel("Frequency (Hz)")
    plt.ylabel("Power")
    plt.grid()
    plt.show()